"""
@filename: python_01.py
@author: Lris
"""
import os
import cv2 as cv
import numpy as np
from skimage.metrics import structural_similarity as ssim

def calculate_ssim_psnr(original, dehazed):
    # 计算 SSIM
    original_gray = cv.cvtColor(original, cv.COLOR_BGR2GRAY)
    dehazed_gray = cv.cvtColor(dehazed, cv.COLOR_BGR2GRAY)
    ssim_index = ssim(original_gray, dehazed_gray)

    # 计算 PSNR
    mse = np.mean((original - dehazed) ** 2)
    if mse == 0:
        psnr = float('inf')
    else:
        max_pixel = 255.0
        psnr = 20 * np.log10(max_pixel / np.sqrt(mse))
    return ssim_index, psnr

def zmMinFilterGray(src, r=5):
    '最小值滤波，r是滤波器半径'
    return cv.erode(src, np.ones((2 * r + 1, 2 * r + 1)))

def darkchannel(Image):
    '''计算每个通道中的最小值，输入Image图像，输出最小值img_min'''
    img_min = np.min(Image, axis=2)
    return img_min

def guided_filter(I, p, r, eps):
    '''导向滤波实现'''
    # 将 I 转换为灰度图像，并确保深度为 CV_32F
    I_gray = cv.cvtColor(I.astype(np.float32), cv.COLOR_BGR2GRAY) / 255.0
    mean_I = cv.boxFilter(I_gray, cv.CV_32F, (r, r))
    mean_p = cv.boxFilter(p, cv.CV_32F, (r, r))
    mean_Ip = cv.boxFilter(I_gray * p, cv.CV_32F, (r, r))
    cov_Ip = mean_Ip - mean_I * mean_p

    mean_II = cv.boxFilter(I_gray * I_gray, cv.CV_32F, (r, r))
    var_I = mean_II - mean_I * mean_I

    a = cov_Ip / (var_I + eps)
    b = mean_p - a * mean_I

    mean_a = cv.boxFilter(a, cv.CV_32F, (r, r))
    mean_b = cv.boxFilter(b, cv.CV_32F, (r, r))

    q = mean_a * I_gray + mean_b
    return q

def dehaze(image, omega=0.95, t0=0.1, r=15, eps=1e-3):
    '''去雾主函数'''
    dark = darkchannel(image)
    A = np.max(dark)
    t = 1 - omega * dark / A
    t = guided_filter(image, t, r, eps)
    t = np.maximum(t, t0)
    J = (image - A) / t[:, :, np.newaxis] + A
    return np.clip(J, 0, 1)  # 确保像素值在 [0, 1] 范围内

# 读取图像并去雾
# 在 DCP.py 中
def process_image(input_path, output_dir):
    # 读取图像
    image = cv.imread(input_path)
    if image is None:
        raise ValueError(f"无法读取图像文件: {input_path}")
    # 转换为 float32 并归一化
    image = image.astype(np.float32) / 255.0
    # 处理图像
    processed_image = dehaze(image)
    # 转换回 uint8
    processed_image = (processed_image * 255).astype(np.uint8)


    # 保存处理后的图像
    filename = os.path.basename(input_path)
    output_filename = f"processed_{filename}"
    output_path = os.path.join(output_dir, output_filename)
    cv.imwrite(output_path, processed_image)

    return output_filename
